import gmsh
import numpy as np
import os
import argparse


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--height", '-H', type=float, help='输入平板的长度(mm)')
    parser.add_argument('--width', '-W', type=float, help='输入平板的宽度')
    parser.add_argument('--thick', '-T', type=float, help='输入平板的厚度')
    args = parser.parse_args()
    height = args.height
    width = args.width
    thick = args.thick
    gmsh.initialize()
    try:
        gmsh.model.add("tablet")
        
        # 1. 创建矩形基面
        rect = gmsh.model.occ.addRectangle(0, 0, 0, width, height)  # 返回 (dim, tag) 中的 tag

        # 2. 同步几何模型（关键步骤！）
        gmsh.model.occ.synchronize()

        # 3. 获取面的边界边（现在 rect 已被同步到 Gmsh 模型）
        edges = gmsh.model.getBoundary([(2, rect)], oriented=False)  # 正确获取边

        # 4. 设置 Transfinite 划分（X/Y方向）
        Nx, Ny = 10, 41
        for edge in edges:
            # 根据实际边方向分配 Nx 或 Ny
            gmsh.model.mesh.setTransfiniteCurve(edge[1], Nx if edge[1] in [edges[0][1], edges[2][1]] else Ny)

        # 5. 设置面为 Transfinite 并重组为四边形
        gmsh.model.mesh.setTransfiniteSurface(rect)
        gmsh.model.mesh.setRecombine(2, rect)

        # 6. 沿Z方向拉伸（1层六面体）
        extrude = gmsh.model.occ.extrude(
            [(2, rect)], 0, 0, thick,
            numElements=[1],  # Z方向1层
            recombine=True    # 生成六面体
        )
        gmsh.model.occ.synchronize()

        # 7. 添加物理组并生成网格
        tablet_volumes = [tag for (dim, tag) in extrude if dim == 3]
        gmsh.model.addPhysicalGroup(3, tablet_volumes, name="tablet")
        gmsh.model.mesh.generate(3)

        output_dir = os.path.join(os.getcwd(), "data", "geo")
        filename = 'tablet_' + 'w' + str(round(width, 2)) + '_h' + str(round(height, 2)) + '_t' + str(round(thick, 2)) + '.msh'
        filename = os.path.join(output_dir, filename)
        gmsh.write(filename)
        gmsh.fltk.run()  # 可视化
        gmsh.clear()  # 清空当前模型
    finally:
        gmsh.finalize()

    print("Generated tablet.geo_unrolled and tablet.msh")

if __name__ == "__main__":
    main()